import pandas as pd
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import f1_score
from sklearn.model_selection import train_test_split

def load_data(train_path, test_path):
    train_data = pd.read_json(train_path, lines=True)
    test_data = pd.read_json(test_path, lines=True)
    return train_data, test_data

def preprocess_and_extract_features(train_data, test_data):
    vectorizer = TfidfVectorizer(max_features=5000)
    X_train = vectorizer.fit_transform(train_data['text'])
    X_test = vectorizer.transform(test_data['text'])
    y_train = train_data['label']
    return X_train, X_test, y_train

def train_model(X_train, y_train):
    model = LogisticRegression()
    model.fit(X_train, y_train) 
    return model

def evaluate_model(model, X_train, y_train):
    y_pred = model.predict(X_train)
    f1 = f1_score(y_train, y_pred)
    print(f'训练集的F1分数为: {f1}')

def predict_on_test_data(model, X_test):
    test_predictions = model.predict(X_test)
    return test_predictions

def save_predictions(predictions, output_path):
    with open(output_path, 'w') as f:
        for prediction in predictions:
            f.write(f'{prediction}\n')

def main(train_path, test_path, output_path):
    
    train_data, test_data = load_data(train_path, test_path)

    X_train, X_test, y_train = preprocess_and_extract_features(train_data, test_data)

    model = train_model(X_train, y_train)

    evaluate_model(model, X_train, y_train)

    test_predictions = predict_on_test_data(model, X_test)

    save_predictions(test_predictions, output_path)

if __name__ == '__main__':
    train_path = 'train.jsonl'
    test_path = 'test.jsonl'
    output_path = 'predictions.txt'  
    main(train_path, test_path, output_path)
